import org.apache.hadoop.fs.FSDataInputStream;
import org.apache.hadoop.fs.FileSystem;
import org.apache.hadoop.fs.Path;
import org.apache.hadoop.io.DataOutputBuffer;
import org.apache.hadoop.io.Text;
import org.apache.hadoop.io.Writable;
import org.apache.hadoop.io.WritableComparable;
import org.apache.hadoop.io.LongWritable;
import org.apache.hadoop.mapred.*;

import java.io.IOException;

/**
 * This class defines a new input format that's aware of the structutre of an
 * XML page.
 *
 * It was originally posted org.apache.hadoop.core-user by Colin Evans
 * (colin@metaweb.com). A copy of that posting can be found at
 * http://markmail.org/message/doebi7mst4mh4ohj as of 30 November 2008.
 *
 * We modified the class as posted to work with in the Hadoop 17.2.1
 * framework, as well as be specific to our task. Since we didn't write this
 * class, and it's at heart a simple extension of TextInputFormat, comments
 * and documentation below are sparse indeed.
 *
 * @author Colin Evans
 * @author Ian Voysey
 * @author Ryan Hofler
 */
public class XmlInputFormat extends TextInputFormat {
    public static final String START_TAG_KEY = "<page>";
    public static final String END_TAG_KEY = "</page>";

    public void configure(JobConf jobConf) {
        super.configure(jobConf);
    }

    public RecordReader<LongWritable, Text> getRecordReader(InputSplit inputSplit,
							    JobConf jobConf,
							    Reporter reporter)
	throws IOException {
        return new XmlRecordReader<LongWritable, Text>((FileSplit) inputSplit, jobConf);
    }

    public static class XmlRecordReader<WritableComparable, Writable>
	implements RecordReader<WritableComparable, Writable> {
        private byte[] startTag;
        private byte[] endTag;
        private long start;
        private long end;
        private FSDataInputStream fsin;
        private DataOutputBuffer buffer = new DataOutputBuffer();

        public XmlRecordReader(FileSplit split, JobConf jobConf) throws IOException {
            startTag = ("<page>").getBytes("utf-8");
            endTag = ("</page>").getBytes("utf-8");

            // open the file and seek to the start of the split
            start = split.getStart();
            end = start + split.getLength();
            Path file = split.getPath();
            FileSystem fs = file.getFileSystem(jobConf);
            fsin = fs.open(split.getPath());
            fsin.seek(start);
        }

        public boolean next(WritableComparable key, Writable value) throws IOException {
            if (fsin.getPos() < end) {
                if (readUntilMatch(startTag, false)) {
                    try {
                        buffer.write(startTag);
                        if (readUntilMatch(endTag, true)) {
                            ((Text) key).set(Long.toString(fsin.getPos()));
                            ((Text) value).set(buffer.getData(), 0, buffer.getLength());
                            return true;
                        }
                    }
                    finally {
                        buffer.reset();
                    }
                }
            }
            return false;
        }


	// this is a valid type cast; i have no idea why java isn't okay with
	// it, and it works
        @SuppressWarnings("unchecked")
        public WritableComparable createKey() {
            return (WritableComparable)(new Text());
        }

	// similar
        @SuppressWarnings("unchecked")
        public Writable createValue() {
            return (Writable)(new Text());
        }

        public long getPos() throws IOException {
            return fsin.getPos();
        }

        public void close() throws IOException {
            fsin.close();
        }

        public float getProgress() throws IOException {
            return ((float) (fsin.getPos() - start)) / ((float) (end - start));
        }

        private boolean readUntilMatch(byte[] match, boolean withinBlock) throws IOException {
            int i = 0;
            while (true) {
                int b = fsin.read();
                // end of file:
                if (b == -1) return false;
                // save to buffer:
                if (withinBlock) buffer.write(b);

                // check if we're matching:
                if (b == match[i]) {
                    i++;
                    if (i >= match.length) return true;
                } else i = 0;
                // see if we've passed the stop point:
                if(!withinBlock && i == 0 && fsin.getPos() >= end) return false;
            }
        }
    }
}
